# _*_ coding: utf-8 _*_
import os
import pickle
from random import sample, shuffle, uniform, randint

import numpy as np
from tqdm import trange

from Encodings import MultiToursSolution


def copy(target):
    return pickle.loads(pickle.dumps(target))


class MA:
    def __init__(self, data):
        self.name = __class__.__name__
        # -- read instance from instance data --
        self.data = copy(data)
        self.V = self.data.cities  # city list
        self.n = self.data.n  # number of cities
        self.m = self.data.m  # number of salesmen
        self.A = self.data.A  # accessibility matrix
        self.Depot = self.data.depot  # depot
        self.V.remove(self.Depot)  # exclude depot
        self.Distance = self.data.distance_matrix  # distance matrix
        self.tour_cost = self.data.cost  # cost function for a single tour
        self.access = self.data.access  # fast computation access func
        # -- parameters ---
        self.NP = 20  # number of population
        self.NS = 10  # number of offspring
        self.Pr = 1.0
        self.limit_stagnation = 20  # Ls
        # -- solution initialization --
        self.Population = [MultiToursSolution().init_solution(cities=self.V, distance=self.Distance, access=self.access,
                                                              depot=self.Depot, m=self.m)
                           for _ in range(self.NP)]
        self.optima = copy(min(self.Population, key=lambda x: x.total))
        self.evo_list = []
        # -- auxiliary --
        self.entropy = []  # assignment entropy to measure the diversity of population
        self.n_unique = []  # the number of unique solutions in the current population
        self.n_children_avg = []  # the average number of children that an individual can produce
        self.n_update = []  # the number of updated pA in a generation
        self.cross = None
        self.strategy = None
        self.assignment = None  # the assignment of current population
    
    def optimize(self, iteration=1000, trial=1):
        pbar = trange(iteration)
        pbar.set_description(f'{self.name} on {self.data.name} trial {trial}')
        for _ in pbar:
            # self.assignment = self.count_assignment(self.Population)  # calculate the assignment of current population
            # H = -np.sum(self.assignment / self.NP * np.log(self.assignment / self.NP + 1e-5))  # calculate entropy
            # self.entropy.append(H)
            
            # self.contrastive_cross()
            
            # # record info
            if self.optima.total > min(self.Population, key=lambda _: _.total).total:
                self.optima = copy(min(self.Population, key=lambda _: _.total))
            self.evo_list.append(self.optima.total)
            self.n_unique.append(len(set([p.total for p in self.Population])))
            
            pbar.set_postfix_str(f'`best: {self.optima.total:.2f}')
            
            plt.subplot(1, 2, 1)
            plt.plot(range(len(self.evo_list)), self.evo_list, label=f'cost: {self.evo_list[-1]:.2f}')
            plt.title(f'optimal cost')
            plt.legend()
            plt.subplot(1, 2, 2)
            print(self.optima)
            self.data.draw(plt, routes=self.optima.tours)
            plt.title(f'optimal cost')
            plt.legend()
            
            # plt.clf()
            # plt.subplot(1, 5, 1)
            # plt.plot(range(len(self.evo_list)), self.evo_list, label=f'cost: {self.evo_list[-1]:.2f}')
            # plt.title(f'optimal cost')
            # plt.legend()
            # plt.subplot(1, 5, 2)
            # plt.plot(range(len(self.entropy)), self.entropy, label=f'entropy: {self.entropy[-1]:.2f}')
            # plt.title(f'assignment entropy of population')
            # plt.legend()
            # plt.subplot(1, 5, 3)
            # plt.plot(range(len(self.n_unique)), self.n_unique, label=f'n_unique: {self.n_unique[-1]}')
            # plt.title(f'the number of unique individuals')
            # plt.legend()
            # plt.subplot(1, 5, 4)
            # plt.plot(range(len(self.n_children_avg)), self.n_children_avg,
            #          label=f'n_children_avg: {self.n_children_avg[-1]:.2f}')
            # plt.plot(range(len(self.n_update)), self.n_update, label=f'n_update: {self.n_update[-1]}')
            # plt.title(f'n_children_avg & n_update')
            # plt.legend()
            # plt.subplot(1, 5, 5)
            # stags = [_.stags for _ in self.Population]
            # plt.hist(x=stags, bins=len(stags), range=[0, 20], rwidth=0.8)
            # plt.title(f'stagnation')
            # plt.tight_layout()
            
            plt.pause(0.001)
            plt.show()
        
        plt.show()
        
        return self
    
    def contrastive_cross(self):
        shuffle(self.Population)
        n_children, update = 0, 0
        for i in range(self.NP):
            # pair-wise cross
            pA = self.Population[i % self.NP]
            pB = self.Population[(i + 1) % self.NP]
            pA_pB = copy(pA)
            
            diff = []  # removed cities
            for k in range(len(pA_pB.tours)):  # compare each route
                _diff = set(pA_pB.tours[k]) - set(pB.tours[k])  # cities in pA but not in pB
                for city in _diff:
                    if uniform(0, 1) < self.Pr:  # remove by Pr
                        diff.append(city)
                        pA_pB.tours[k].remove(city)
            
            # generate no-repeating children based on pA
            children = []
            repeat = 0
            while repeat < self.NS and len(children) < self.NS:  # repeating limits
                no_repeating_costs = [child.total for child in children]
                child = self.best_insert(base=pA_pB, diff=diff)
                # child = self.local_search(s=child)
                child.calc_cost(cost_func=self.tour_cost)
                if child.total not in no_repeating_costs:
                    children.append(child)
                else:
                    child = self.rand_insert(base=pA_pB, diff=diff)
                    if child.total not in no_repeating_costs:
                        children.append(child)
                    else:
                        repeat += 1
            n_children += len(children)  # the number of children of pA
            
            children = [child for child in children if child.total < pA.total]  # preserver the improved ones
            fitness = []
            for child in children:
                new_assignment = self.assignment_change(self.assignment, re=pA, ad=child)
                new_entropy = -np.sum(new_assignment / self.NP * np.log(new_assignment / self.NP + 1e-5))
                entropy_loss = self.entropy[-1] - new_entropy
                if entropy_loss <= 0:  # the child improves with no entropy loss at all
                    fitness.append(pA.total - child.total)  # improved total cost
                else:  # the child improves at the cost of entropy loss
                    fitness.append((pA.total - child.total) / entropy_loss)  # improvement on per unit of entropy loss
            
            if children:  # if pA ever has improved children
                child = children[fitness.index(max(fitness))]
                self.Population[i] = copy(child)  # replace pA with the best child
                self.Population[i].stags = 0
                update += 1
            else:
                self.Population[i].stags += 1
                if self.Population[i].stags > self.limit_stagnation:
                    self.Population[i] = self.outer_mutate(self.Population[i])
                    self.Population[i] = self.inner_mutate(self.Population[i])
                    self.Population[i].calc_cost(cost_func=self.tour_cost)
                    self.Population[i].stags = 0
        
        self.n_children_avg.append(n_children / self.NP)
        self.n_update.append(update)
    
    def best_insert(self, base, diff):
        base = copy(base)
        shuffle(diff)
        for city in diff:
            self.optimal_insert(base, city)
        return base
    
    def rand_insert(self, base, diff):
        base = copy(base)
        shuffle(diff)
        for city in diff:
            k = sample(self.salesmen_of(city), 1)[0]
            self.optimal_insert(base, city, k)
        return base
    
    def relocate(self, s):
        for _ in range(self.m):
            k = randint(0, len(s.tours) - 1)
            i = randint(1, len(s.tours[k]) - 2)  # a random city in a random tour
            city = s.tours[k][i]
            s.tours[k].pop(i)  # remove the city
            k = sample(self.salesmen_of(city), 1)[0]  # random choose a feasible tour
            i = randint(1, len(s.tours[k]) - 1)  # insert to a random position
            s.tours[k].insert(i, city)
        return s
    
    def outer_mutate(self, s):
        for k in range(self.m):
            i = randint(1, len(s.tours[k]) - 2)  # choose a random city
            city = s.tours[k][i]
            s.tours[k].pop(i)  # remove the city
            _k = sample(self.salesmen_of(city), 1)[0]  # random choose a feasible tour
            self.optimal_insert(solution=s, city=city, salesman=_k)
        return s
    
    def inner_mutate(self, s):
        for k in range(self.m):
            i = randint(1, len(s.tours[k]) - 2)  # choose a random city
            city = s.tours[k][i]
            s.tours[k].pop(i)  # remove the city
            self.optimal_insert(solution=s, city=city, salesman=k)  # re-insert the city into the same tour
        return s
    
    def local_search(self, s):
        for _ in range(len(s.tours)):
            s.tours[_] = self.two_opt(s.tours[_])
        return s
    
    def two_opt(self, tour):
        for _ in range(self.Opt_depth):
            _i, _j, improve = -1, -1, -1  # best 2-opt move
            for i in range(1, len(tour)):  # find the best 2-opt move
                for j in range(i + 1, len(tour)):
                    d1 = self.Distance[tour[i - 1]][tour[i]] + self.Distance[tour[j - 1]][tour[j]]  # remove edges
                    d2 = self.Distance[tour[i - 1]][tour[j - 1]] + self.Distance[tour[i]][tour[j]]  # insert edges
                    if d1 - d2 > improve:
                        _i, _j, improve = i, j, d1 - d2
            if improve > 0:  # perform the best 2-opt move
                tour = tour[:_i] + list(reversed(tour[_i:_j])) + tour[_j:]
            else:
                break  # no improving edge exchange exits, break
        return tour
    
    def count_assignment(self, pop):
        assigns = np.zeros(shape=(self.m, self.n + 1), dtype=int)  # [m][n]: salesman m is assigned with city n
        for s in pop:
            for k, tour in enumerate(s.tours):
                for city in tour:
                    if city != self.Depot:
                        assigns[k][city] += 1  # the number of individuals that assign city to k
        return assigns
    
    def assignment_change(self, assigns, re, ad):
        assigns = copy(assigns)
        for k, tour in enumerate(re.tours):
            for city in tour:
                if city != self.Depot:
                    assigns[k][city] -= 1
        for k, tour in enumerate(ad.tours):
            for city in tour:
                if city != self.Depot:
                    assigns[k][city] += 1
        return assigns
    
    def assignment_entropy(self, pop):
        assigns = self.count_assignment(pop)
        assigns = assigns / self.NP * np.log(assigns / self.NP + 1e-5)
        return -np.sum(assigns)
    
    @staticmethod
    def average_length(pop):
        return sum([s.total for s in pop]) / len(pop)


def make_dir(path):
    path = path.strip().rstrip("\\")
    isExists = os.path.exists(path)
    if not isExists:
        os.makedirs(path)
        print(path + ' created.')
        return True
    else:
        print(path + ' already exists.')
        return False


def save_tours(path: str, seqs):
    for k, seq in enumerate(seqs):
        mode = 'w' if k == 0 else 'a'
        seq = np.array(seq).reshape(1, len(seq))
        with open(path, mode) as file:
            np.savetxt(file, seq, fmt='%d')


from VCMTSP import Instance
import matplotlib.pyplot as plt

plt.figure(figsize=(10, 4))
case, m = 'att48', 2
data = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
optimizer = MA(data=data)
optimizer.optimize(iteration=1000)

set_IIa = ['kroA100', 'kroA150', 'kroA200'], [2, 4, 6, 8]
set_IIb = ['kroB100', 'kroB150', 'kroB200'], [2, 4, 6, 8]
set_IIc = ['kroC100', 'kroD100', 'kroE100'], [2, 4, 6, 8]

set_I = ['att48', 'eil51', 'berlin52', 'eil76', 'eil101'], [2, 4, 6, 8]
set_II = ['gr202', 'gr229', 'gr431', 'gr666'], [10, 20, 30, 40]
set_III = ['dsj1000', 'pr1002', 'pcb1173'], [40, 60, 80, 100]

test_suite = {'set_I': set_I,
              'set_II': set_II,
              'set_III': set_III}

# if __name__ == '__main__':
#     save_path = rf'./results/{strftime("%Y-%m-%d", localtime())}'
#     make_dir(save_path)
#     parser = argparse.ArgumentParser()
#     parser.add_argument('-S', required=True, type=str, help='The testing dataset!')
#     parser.add_argument('-G', required=True, type=int, help='The max generation!')
#     parser.add_argument('-T', required=True, type=int, help='The times of experiments!')
#     args = parser.parse_args()
#     benchmark = test_suite[f'{args.S}']
#     total_trial = args.T
#     plt.figure(figsize=(4, 3))
#     for case in benchmark[0]:
#         for m in benchmark[1]:
#             instance = Instance(file=f'../../data/TSPLIB_VC/{case}.tsp/{case}.tsp.txt', m=m)
#             n_cities = instance.n
#             n_salesmen = instance.m
#             save_name = f'MA-{case}-m{m}'
#             evo_list, entropy, n_unique, n_update = [], [], [], []
#             tours, total = None, 1e99
#             for trial in range(total_trial):
#                 optimizer = MA(data=instance)
#                 optimizer.NP = 20
#                 optimizer.optimize(iteration=args.G, trial=trial)
#                 evo_list.append(optimizer.evo_list)
#                 entropy.append(optimizer.entropy)
#                 n_unique.append(optimizer.n_unique)
#                 n_update.append(optimizer.n_update)
#                 if optimizer.Best.total < total:
#                     tours, total = optimizer.Best.tours, optimizer.Best.total
#
#             np.savetxt(f'{save_path}/{save_name}-evo_list.txt', np.array(evo_list).data, fmt='%.2f')
#             np.savetxt(f'{save_path}/{save_name}-entropy.txt', np.array(entropy).data, fmt='%.2f')
#             np.savetxt(f'{save_path}/{save_name}-n_unique.txt', np.array(n_unique).data, fmt='%.2f')
#             np.savetxt(f'{save_path}/{save_name}-n_update.txt', np.array(n_update).data, fmt='%.2f')
#             save_tours(f'{save_path}/{save_name}-tours.txt', tours)
